import asyncio
from typing import Optional
from contextlib import AsyncExitStack

from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

from anthropic import Anthropic
from anthropic.types import ContentBlock, TextBlock, ToolUseBlock
from dotenv import load_dotenv
# Import Ollama types for response parsing

from ollama import Client, AsyncClient
from ollama._types import ChatResponse

import logging
import os, sys
import json
from typing import Dict, List, Any, Optional, Union, TypeVar
from typing_extensions import Literal

JsonDict = Dict[str, Any]
MessageType = Dict[str, Any]
OllamaToolType = Dict[str, Any]
EMPTY_JSON_VALUES = ('', '{}')
T = TypeVar('T')

load_dotenv()  # load environment variables from .env

BASE_URL = "http://localhost:11434"

current_path = os.path.dirname(os.path.realpath(__file__))

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s')
logger = logging.getLogger(__name__)

class Message:
    id: str
    """Unique object identifier.

    The format and length of IDs may change over time.
    """
    content: List[ContentBlock]

    role: Literal["assistant"]
    """Conversational role of the generated message.

    This will always be `"assistant"`.
    """
    model: str

    def __init__(self, role:str, model:str):
        self.role = role
        self.content = []
        self.model = model
        pass

# Reference: git@github.com:modelcontextprotocol/quickstart-resources.git 
class MCPClient:
    def __init__(self):
        # Initialize session and client objects
        self.session: Optional[ClientSession] = None
        self.exit_stack = AsyncExitStack()
        #self.anthropic = Anthropic()
        #self.anthropic = Anthropic(base_url=BASE_URL,  api_key="fake-key", auth_token="fake-token")
        self.ollama = AsyncClient(host=BASE_URL)

    async def connect_to_server(self, server_script_path: str):
        """Connect to an MCP server
        
        Args:
            server_script_path: Path to the server script (.py or .js)
        """
        is_python = server_script_path.endswith('.py')
        is_js = server_script_path.endswith('.js')
        if not (is_python or is_js):
            raise ValueError("Server script must be a .py or .js file")
            
        command = "python" if is_python else "node"
        server_params = StdioServerParameters(
            command=command,
            args=[server_script_path],
            env=None
        )
        stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
        self.stdio, self.write = stdio_transport
        self.session = await self.exit_stack.enter_async_context(ClientSession(self.stdio, self.write))
        
        await self.session.initialize()
        
        # List available tools
        response = await self.session.list_tools()
        tools = response.tools
        print("\nConnected to server with tools:", [tool.name for tool in tools])

    def convert_to_anthropic_format(self, response: ChatResponse) -> Message:
        # anthropic.types.Message must have role=assistant in respone
        message=Message(role='assistant', model=response.model)
        if response.message.tool_calls is None:
            content: ContentBlock= TextBlock(
                text=response.message.content,
                type="text"
            )
            message.content.append(content)
            logger.debug(f"response.message: {response.message.content}")
        else:
            for idx, tc in enumerate(response.message.tool_calls):
                content: ContentBlock= ToolUseBlock(
                    id = f'{idx}',
                    name=tc.function.name,
                    input=tc.function.arguments,
                    type="tool_use"
                )
                message.content.append(content)
        return message
    def convert_mcp_tool(self, idx: int, tool: Dict[str, Any]) -> Optional[OllamaToolType]:
        """
        Convert a single MCP tool definition into Ollama tool format.
        
        Args:
            idx: Index for logging purposes
            tool: MCP tool definition
            
        Returns:
            Tool in Ollama format or None if processing fails
        """
        try:
            # Store the original name in our mapping
            original_name = tool["name"]
            logger.debug(f"Processing tool [{idx}]: {original_name}")
            # Get parameter properties
            # properties, required = extract_tool_parameters(tool)
            properties = {}
            required = []
            if "parameters" in tool:
                if isinstance(tool["parameters"], dict):
                    properties = tool["parameters"].get("properties", {})
                    required = tool["parameters"].get("required", [])
                    logger.debug(f"Tool {original_name} has parameters: properties={list(properties.keys())}, required={required}")
                else:
                    logger.warning(f"Tool {original_name} parameters not a dict: {type(tool['parameters'])}")
            else:
                logger.debug(f"Tool {original_name} has no parameters defined")
 
            # Create tool in Ollama's expected format based on docs
            ollama_tool = {
                "type": "function",
                "function": {
                    "name": original_name,
                    "description": tool.get("description", ""),
                    "parameters": {
                        "type": "object",
                        "properties": properties,
                        "required": required
                    }
                }
            }
            logger.debug(f"Convert tool {original_name} to Ollama format successfully")
            return ollama_tool
        except Exception as e:
            logger.error(f"Error processing tool: {e}")
            return None
    def convert_mcp_tools(self, mcp_tools: Union[List[Any]]) -> List[OllamaToolType]:
        """
        Convert MCP tools format to Ollama tool format according to the Ollama SDK docs.
        
        Args:
            mcp_tools: Tools in MCP format (list, dict with 'tools' key, or object with 'tools' attribute)
            
        Returns:
            List of tools formatted for Ollama's API
        """
        logger.debug(f"Converting {len(mcp_tools)} MCP tools to Ollama format")
        # Extract tools from the input based on its type
        logger.debug(f"origin mcp tools: {json.dumps(mcp_tools, indent=2)}")
        ollama_tools = []
        # Process each tool in the list
        for idx, tool in enumerate(mcp_tools):
            if "name" in tool and "description" in tool:
                # Process valid tool
                ollama_tool = self.convert_mcp_tool(idx, tool)
                if ollama_tool:
                    ollama_tools.append(ollama_tool)
            else:
                logger.warning(f"Tool missing required attributes: has name = {'name' in tool}, has description = {'description' in tool}")
 
        return ollama_tools 
    async def process_query(self, query: str) -> str:
        """Process a query using Claude and available tools"""
        """
        messages:List[MessageType]= [
            {
                "role": "system", 
                "content": "You are a weather forecast assistant with ability to use tool function call according to location or geographical coordinates."
            },
            {
                "role": "user",
                "content": query
            }
        ]
        """
        messages:List[MessageType]= [ 
            {
                "role": "system", 
                "content": "You are a weather forecast assistant with ability to use tool function call according to location or geographical coordinates."
            }]
        messages.append(
            {
                "role": "user",
                "content": query
            }
        )

        response = await self.session.list_tools()
        logger.debug(f"MCP original tools: {response.tools}")
        available_tools = [{ 
            "name": tool.name,
            "description": tool.description,
            "parameters": tool.inputSchema
            #"input_schema": tool.inputSchema
        } for tool in response.tools]

        # Initial Claude API call
        """
        response = self.anthropic.messages.create(
            #model="claude-3-5-sonnet-20241022",
            model="qwen3:0.6b",
            max_tokens=1000,
            messages=messages,
            tools=available_tools
        )
        """
       
        logger.info(f"preprocess messages: {messages}")
        ollama_tools=self.convert_mcp_tools(available_tools)
        logger.debug(f"Ollama format tools: {ollama_tools}")
        response: ChatResponse = await self.ollama.chat(
            model="qwen3:0.6b",
            messages=messages,
            tools=ollama_tools,
            stream=False,
        )
        logger.info(f"ollama response: {response}")
        anthropic_response=self.convert_to_anthropic_format(response)

        # Process response and handle tool calls
        final_text = []

        logger.debug(f"anthropic response: {anthropic_response.model}")
        #for content in response.content:
        for content in anthropic_response.content:
            if content.type == 'text':
                final_text.append(content.text)
            elif content.type == 'tool_use':
                tool_name = content.name
                tool_args = content.input
                
                # Execute tool call
                result = await self.session.call_tool(tool_name, tool_args)
                final_text.append(f"[Calling tool {tool_name} with args {tool_args}]")
                logger.info(f"[Calling tool {tool_name} with args {tool_args}]: result={result}")
                # Continue conversation with tool results
                if hasattr(content, 'text') and content.text:
                    messages.append({
                      "role": "assistant",
                      "content": content.text
                    })
                messages.append({
                    "role": "user", 
                    "content": result.content[0].text
                })

                """
                # Get next response from Claude
                response = self.anthropic.messages.create(
                    model="claude-3-5-sonnet-20241022",
                    max_tokens=1000,
                    messages=messages,
                )
                """
                response: ChatResponse = await self.ollama.chat(
                    model="qwen3:0.6b",
                    messages=messages,
                    tools=None,
                    stream=False,
                )
                #final_text.append(response.content[0].text)
                final_text.append(response.message.content)

        return "\n".join(final_text)

    async def chat_loop(self):
        """Run an interactive chat loop"""
        print("\nMCP Client Started!")
        print("Type your queries or 'quit' to exit.")
        
        while True:
            try:
                query = input("\nQuery: ").strip()
                
                if query.lower() == 'quit':
                    break
                    
                response = await self.process_query(query)
                print("\n" + response)
                    
            except Exception as e:
                print(f"\nError: {str(e)}")
    
    async def cleanup(self):
        """Clean up resources"""
        await self.exit_stack.aclose()

async def main():
    if len(sys.argv) < 2:
        print("Usage: python client.py <path_to_server_script>")
        sys.exit(1)

    client = MCPClient()
    try:
        await client.connect_to_server(sys.argv[1])
        #await client.start()
        
        await client.chat_loop()
    finally:
        await client.cleanup()

if __name__ == "__main__":
    asyncio.run(main())
